Skip to content

fix(datasets): backfill .static on caches predating it (PF test/predict) - #125

Merged
albanpuech merged 2 commits into
mainfrom
fix/pf-static-backfill
Sep 30, 2026
Merged

fix(datasets): backfill .static on caches predating it (PF test/predict)#125
albanpuech merged 2 commits into
mainfrom
fix/pf-static-backfill

Conversation

@romeokienzler

@romeokienzler romeokienzler commented Sep 24, 2026

Copy link
Copy Markdown
Collaborator

Problem

The PF test/predict path crashes on any dataset processed before .static was introduced:

File "gridfm_graphkit/tasks/pf_task.py", line 132, in test_step
    RemovePFMask()(batch)
File "gridfm_graphkit/datasets/masking.py", line 226, in forward
    ].static
AttributeError: 'NodeStorage' object has no attribute 'static'

RemovePFMask (in test_step / predict_step) restores static limit columns from the .static
snapshot that build_hetero_data attaches at graph-build time. But .static was added after many
datasets were already processed to disk. HeteroGridDatasetDisk.process() skips regeneration when a
processed_raw_files.done marker exists, so those stale caches are never rebuilt — their graphs carry
the limit columns only inside bus.x / branch edge_attr, and loading one crashes the PF
test/predict path.

This was surfaced by running the training integration tests (integrationtests/test_base_set.py::test_train_pf)
on main: the committed case14 fixture ships pre-processed data_index_*.pt graphs built before
.static existed, so every PF run aborts in the test phase. (OPF is unaffected — it doesn't use .static.)

Fix

  • Add backfill_static(data) in datasets/graph_builder.py: if a graph lacks .static, reconstruct it
    from the raw limit columns still present in bus.x / branch edge_attr — identical to what
    build_hetero_data stores. Freshly-processed graphs already carry .static, so it's a no-op for them.
  • Invoke it in HeteroGridDatasetDisk.get() right after HeteroData.from_dict and before
    normalisation, so the captured snapshot matches build-time raw semantics (and branch masking slices
    it consistently afterward).
  • Centralise the limit-column indices as BUS_STATIC_COLS / BRANCH_STATIC_COLS so the build path and
    the backfill path share one source of truth.

No dataset regeneration required — stale caches become forward-compatible on load.

Validation

On vela (torch 2.13.0+cu130, H100), pytest integrationtests against the stale case14 fixture:
2 passed in 0:45:40 (PF + OPF), no AttributeError. Before the fix, PF crashed in the test phase.

🤖 Generated with Claude Code

@romeokienzler @claude
fix(datasets): backfill .static on caches predating it (PF test/predict)
e09dcc3 
RemovePFMask (PF test_step/predict_step) reads `.static` limit snapshots that
build_hetero_data attaches at graph-build time. Datasets processed before
`.static` was introduced (PRs #99/#105/#122) carry the limit columns only inside
`bus.x` / branch `edge_attr`, so loading such a cache and running the PF
test/predict path crashes with:

    AttributeError: 'NodeStorage' object has no attribute 'static'

process() skips regeneration when a `processed_raw_files.done` marker exists, so
stale caches are never rebuilt. Add backfill_static(), invoked in
HeteroGridDatasetDisk.get() right after from_dict and before normalisation, which
reconstructs `.static` from the raw (pre-mask, pre-normalisation) limit columns —
identical to what build_hetero_data stores. Freshly processed graphs already
carry `.static`, so the backfill is a no-op for them.

Also centralise the limit-column indices as BUS_STATIC_COLS / BRANCH_STATIC_COLS
so build_hetero_data and backfill_static share one source of truth.

Verified on vela (torch 2.13.0, H100): integrationtests PF + OPF both pass
against the stale case14 fixture (2 passed, 0:45:40); no AttributeError.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Romeo Kienzler <romeo.kienzler1@ibm.com>
@romeokienzler

romeokienzler commented Sep 24, 2026

Copy link
Copy Markdown
Collaborator Author

@romeokienzler thanks for the fix — the diagnosis is clear and the BUS_STATIC_COLS/BRANCH_STATIC_COLS single-source-of-truth refactor is a nice touch. The placement of backfill_static in get() (after from_dict, before normalisation) matches the build-time raw semantics.

What's needed

  • Unit test for backfill_static. The integration run on vela is great validation, but per the contribution checklist new functionality needs a unit test in CI too. tests/test_remove_pf_mask.py already builds a graph with .static; a small companion test could (a) delete .static from a built graph and assert backfill_static reconstructs it to equal the original snapshot, and (b) assert it's a no-op when .static is already present. That locks in the exact behaviour you verified.
  • CI still pending — pytests / pre-commit / CodeQL haven't reported yet; worth a glance once they land. DCO is signed off ✅. Note the repo-wide pip-audit failure (PYSEC-2026-3624 / lightning) is a known infra blocker being fixed separately — not this PR's fault.

Otherwise this looks tightly scoped and forward-compatible with no reprocessing required, which is the right call. Deferring the merge decision to a maintainer.

— 🤖 _automated pre-review; a maintainer will follow up_

@romeokienzler romeokienzler mentioned this pull request Sep 24, 2026
Merged
@romeokienzler

romeokienzler commented Sep 30, 2026

Copy link
Copy Markdown
Collaborator Author

@romeokienzler the merge of latest main (0be01d0) synced the branch — no more BEHIND/rebase concern, and I see a maintainer has already approved. DCO ✅, and the security scans that have reported (Trivy, Bandit, detect-secrets) are green.

Two quick notes on where things stand vs. my earlier review:

  • Unit test still the one open item. The merge only pulled in main; the code diff is unchanged, so backfill_static still has no CI-level unit test. The vela integration run is solid validation, but per the contribution checklist a small unit test in tests/ — (a) delete .static from a built graph and assert backfill_static reconstructs the original snapshot, (b) assert it's a no-op when .static is already present — would lock in the behaviour and guard against regressions.
  • CI still settling — pytests / pre-commit / CodeQL / security are pending; worth a glance once they land. The pending pip-audit is the known repo-wide PYSEC-2026-3624 / lightning blocker, not this PR's fault.

Otherwise unchanged from before: tightly scoped, forward-compatible, no reprocessing needed. Deferring the merge decision to the approving maintainer.

— 🤖 _automated pre-review; a maintainer will follow up_

@albanpuech
albanpuech merged commit 1e370e5 into main Sep 30, 2026
11 checks passed
@albanpuech
albanpuech deleted the fix/pf-static-backfill branch September 30, 2026 08:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@albanpuech albanpuech albanpuech approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@romeokienzler @albanpuech